#!/usr/bin/env python3
"""
基于 MCP 协议的多代理文章自动编写系统
主入口文件
"""

import os
import json
import time
from datetime import datetime
from typing import Dict, Any

from agents.research_agent import ResearchAgent
from agents.writing_agent import WritingAgent
from agents.review_agent import ReviewAgent
from agents.polishing_agent import PolishingAgent
from tools.file_tools import FileTools
import config

class MultiAgentWritingSystem:
    """多代理文章编写系统"""

    def __init__(self):
        self.agents = {
            "research": ResearchAgent(),
            "writing": WritingAgent(),
            "review": ReviewAgent(),
            "polishing": PolishingAgent()
        }
        self.file_tools = FileTools()
        self.execution_log = {
            "start_time": "",
            "end_time": "",
            "user_query": "",
            "agents_execution": [],
            "final_result": {},
            "error_log": [],
            "retry_log": []
        }

        # 创建输出目录
        if not os.path.exists("output"):
            os.makedirs("output")

    def run(self, user_query: str, writing_style: str = "professional",
            target_length: str = "中等") -> Dict[str, Any]:
        """运行多代理系统"""
        print("🚀 启动多代理文章编写系统...")
        print(f"📝 用户查询: {user_query}")
        print(f"🎨 写作风格: {writing_style}")
        print(f"📏 文章长度: {target_length}")
        print("-" * 50)

        self.execution_log["start_time"] = self._get_timestamp()
        self.execution_log["user_query"] = user_query

        task_config = {
            "query": user_query,
            "writing_style": writing_style,
            "target_length": target_length
        }

        try:
            # 1. 研究代理
            research_result = self._execute_agent_with_retry(
                "research", self.agents["research"], task_config
            )

            # 2. 撰写代理（依赖研究结果）
            writing_result = self._execute_agent_with_retry(
                "writing", self.agents["writing"], task_config,
                research_data=research_result.get("research_data", "")
            )

            # 3. 审核代理（依赖写作结果）
            review_result = self._execute_agent_with_retry(
                "review", self.agents["review"], task_config,
                article_content=writing_result.get("article_content", "")
            )

            # 4. 润色代理（依赖审核结果）
            polishing_result = self._execute_agent_with_retry(
                "polishing", self.agents["polishing"], task_config,
                article_content=writing_result.get("article_content", ""),
                review_suggestions=review_result.get("suggestions", [])
            )

            # 生成最终结果
            final_result = self._generate_final_output(
                task_config, research_result, writing_result,
                review_result, polishing_result
            )

            self.execution_log["end_time"] = self._get_timestamp()
            self.execution_log["final_result"] = final_result

            # 保存结果
            self._save_results(final_result)

            print("🎉 文章编写完成！")
            print(f"📁 输出文件: output/final_article.md")
            print(f"📊 执行日志: output/execution_log.json")

            return final_result

        except Exception as e:
            error_msg = f"系统执行失败: {str(e)}"
            print(f"❌ {error_msg}")
            self.execution_log["error_log"].append(error_msg)
            raise

    def _execute_agent_with_retry(self, agent_type: str, agent, task_config: Dict[str, Any],
                                 **kwargs) -> Dict[str, Any]:
        """带重试机制的代理执行"""
        max_retries = config.AGENT_CONFIG.get(f"{agent_type}_agent", {}).get("max_retries", 2)

        for attempt in range(max_retries + 1):
            try:
                result = agent.execute(task_config, **kwargs)

                # 记录执行日志
                agent_log = {
                    "agent": agent.name,
                    "attempt": attempt + 1,
                    "status": result.get("status", "unknown"),
                    "timestamp": result.get("timestamp", ""),
                    "details": result
                }
                self.execution_log["agents_execution"].append(agent_log)

                if result.get("status") == "success":
                    return result
                else:
                    print(f"⚠️ {agent.name} 执行失败，准备重试...")
                    self.execution_log["retry_log"].append({
                        "agent": agent.name,
                        "attempt": attempt + 1,
                        "error": result.get("error", "未知错误")
                    })

            except Exception as e:
                print(f"❌ {agent.name} 执行异常: {str(e)}")
                self.execution_log["error_log"].append({
                    "agent": agent.name,
                    "error": str(e),
                    "attempt": attempt + 1
                })

        # 所有重试都失败
        raise Exception(f"{agent.name} 执行失败，已达到最大重试次数")

    def _generate_final_output(self, task_config: Dict[str, Any], research_result: Dict[str, Any],
                              writing_result: Dict[str, Any], review_result: Dict[str, Any],
                              polishing_result: Dict[str, Any]) -> Dict[str, Any]:
        """生成最终输出"""
        final_content = f"""# 最终文章成果

## 文章信息
- **主题**: {task_config['query']}
- **风格**: {task_config['writing_style']}
- **长度**: {task_config['target_length']}
- **生成时间**: {self._get_timestamp()}

## 文章内容

{polishing_result.get('polished_content', '')}

## 质量评估
- **初稿评分**: {review_result.get('overall_score', 0)}/10
- **最终评分**: {polishing_result.get('final_score', 0)}/10
- **改进数量**: {polishing_result.get('improvements_made', 0)} 处

## 代理协作过程
{self._format_agent_process()}

## 异常处理日志
{self._format_error_log()}
"""

        return {
            "final_article": final_content,
            "quality_metrics": {
                "initial_score": review_result.get('overall_score', 0),
                "final_score": polishing_result.get('final_score', 0),
                "improvement": polishing_result.get('final_score', 0) - review_result.get('overall_score', 0)
            },
            "execution_summary": {
                "total_agents": len(self.agents),
                "successful_agents": len([log for log in self.execution_log["agents_execution"]
                                        if log.get("status") == "success"]),
                "total_retries": len(self.execution_log["retry_log"]),
                "total_errors": len(self.execution_log["error_log"])
            }
        }

    def _format_agent_process(self) -> str:
        """格式化代理协作过程"""
        process_text = ""
        for i, log in enumerate(self.execution_log["agents_execution"], 1):
            process_text += f"{i}. {log['agent']} (尝试{log['attempt']}次) - {log['status']}\n"
        return process_text

    def _format_error_log(self) -> str:
        """格式化错误日志"""
        if not self.execution_log["error_log"]:
            return "无异常情况"

        error_text = ""
        for i, error in enumerate(self.execution_log["error_log"], 1):
            error_text += f"{i}. {error}\n"
        return error_text

    def _save_results(self, final_result: Dict[str, Any]):
        """保存结果到文件"""
        # 保存最终文章
        article_filename = self.file_tools.save_article(
            final_result["final_article"], "final_article.md"
        )

        # 保存执行日志
        log_filename = self.file_tools.save_execution_log(self.execution_log)

        print(f"💾 结果已保存: {article_filename}, {log_filename}")

    def _get_timestamp(self) -> str:
        return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def main():
    """主函数"""
    print("=" * 60)
    print("基于 MCP 协议的多代理文章自动编写系统")
    print("=" * 60)

    # 用户输入
    user_query = input("请输入您要写的文章主题: ").strip()
    if not user_query:
        user_query = "AI Agent 技术的发展现状与未来趋势"

    print("\n可选写作风格:")
    for style_key, style_name in config.WRITING_STYLES.items():
        print(f"  {style_key}: {style_name}")

    writing_style = input("请选择写作风格 (默认: professional): ").strip()
    if writing_style not in config.WRITING_STYLES:
        writing_style = "professional"

    target_length = input("请选择文章长度 (短篇/中等/长篇, 默认: 中等): ").strip()
    if target_length not in ["短篇", "中等", "长篇"]:
        target_length = "中等"

    print("\n" + "=" * 60)

    # 运行系统
    system = MultiAgentWritingSystem()

    try:
        result = system.run(user_query, writing_style, target_length)

        # 显示摘要信息
        print("\n📊 执行摘要:")
        print(f"   主题: {user_query}")
        print(f"   风格: {config.WRITING_STYLES[writing_style]}")
        print(f"   长度: {target_length}")
        print(f"   质量评分: {result['quality_metrics']['final_score']}/10")
        print(f"   质量提升: +{result['quality_metrics']['improvement']:.1f}")
        print(f"   成功代理: {result['execution_summary']['successful_agents']}/{result['execution_summary']['total_agents']}")

    except Exception as e:
        print(f"系统运行失败: {e}")

if __name__ == "__main__":
    main()